{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using FalkorGraphRagCapability with agents for GraphRAG Question & Answering\n",
    "\n",
    "AG2 provides GraphRAG integration using agent capabilities. This is an example to integrate FalkorDB (a Knowledge Graph database).\n",
    "\n",
    "````{=mdx}\n",
    ":::info Requirements\n",
    "FalkorDB's GraphRAG-SDK is a dependency for this notebook, which can be installed with ag2 via pip:\n",
    "\n",
    "```bash\n",
    "pip install -U ag2[openai,graph-rag-falkor-db]\n",
    "```\n",
    "\n",
    "> **Note:** If you have been using `autogen` or `ag2`, all you need to do is upgrade it using:  \n",
    "> ```bash\n",
    "> pip install -U autogen[openai,graph-rag-falkor-db]\n",
    "> ```\n",
    "> or  \n",
    "> ```bash\n",
    "> pip install -U ag2[openai,graph-rag-falkor-db]\n",
    "> ```\n",
    "> as `autogen`, and `ag2` are aliases for the same PyPI package.  \n",
    "\n",
    "\n",
    "For more information, please refer to the [installation guide](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/installing-ag2).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set Configuration and OpenAI API Key\n",
    "\n",
    "By default, in order to use FalkorDB you need to have an OpenAI key in your environment variable `OPENAI_API_KEY`.\n",
    "\n",
    "You can utilise an OAI_CONFIG_LIST file and extract the OpenAI API key and put it in the environment, as will be shown in the following cell.\n",
    "\n",
    "Alternatively, you can load the environment variable yourself.\n",
    "\n",
    "````{=mdx}\n",
    ":::tip\n",
    "Learn more about configuring LLMs for agents [here](https://docs.ag2.ai/latest/docs/user-guide/basic-concepts/llm-configuration).\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import autogen\n",
    "\n",
    "llm_config = autogen.LLMConfig.from_json(path=\"OAI_CONFIG_LIST\")\n",
    "\n",
    "# Put the OpenAI API key into the environment\n",
    "os.environ[\"OPENAI_API_KEY\"] = llm_config.config_list[0][\"api_key\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "````{=mdx}\n",
    ":::important\n",
    "The default model for loading graph data and answering questions using FalkorDB's SDK is OpenAI's GPT 4o and this can be changed by setting the `model` parameter on the FalkorGraphQueryEngine.\n",
    ":::\n",
    "````"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a Knowledge Graph with Your Own Data\n",
    "\n",
    "**Note:** You need to have a FalkorDB graph database running. If you are running one in a Docker container, please ensure your Docker network is setup to allow access to it.\n",
    "\n",
    "In this example, the FalkorDB endpoint is set to host=\"172.18.0.3\" and port=6379, please adjust accordingly. For how to set up FalkorDB, please refer to https://docs.falkordb.com/\n",
    "\n",
    "Below, we have some sample data from IMDB on the movie 'The Matrix'. See the [contents of the file here](https://github.com/ag2ai/ag2/blob/main/test/agentchat/contrib/graph_rag/the_matrix.txt).\n",
    "\n",
    "We then initialise the database with that text document, creating the graph in FalkorDB."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A Simple Example\n",
    "\n",
    "In this example, the graph ontology is auto-generated. This allows you to load data without specifying the specific types of entities and relationships that will make up the database (however, this may not be optimal and not cost efficient)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen import ConversableAgent, UserProxyAgent\n",
    "from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType\n",
    "from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine\n",
    "from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability\n",
    "\n",
    "# Auto generate graph schema from unstructured data\n",
    "input_path = \"../test/agentchat/contrib/graph_rag/the_matrix.txt\"\n",
    "input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path)]\n",
    "\n",
    "# Create FalkorGraphQueryEngine\n",
    "query_engine = FalkorGraphQueryEngine(\n",
    "    name=\"The_Matrix_Auto\",\n",
    "    host=\"172.17.0.4\",  # Change\n",
    "    port=6379,  # if needed\n",
    ")\n",
    "\n",
    "# Ingest data and initialize the database\n",
    "query_engine.init_db(input_doc=input_documents)\n",
    "\n",
    "# Create a ConversableAgent (no LLM configuration)\n",
    "graph_rag_agent = ConversableAgent(\n",
    "    name=\"matrix_agent\",\n",
    "    human_input_mode=\"NEVER\",\n",
    ")\n",
    "\n",
    "# Associate the capability with the agent\n",
    "graph_rag_capability = FalkorGraphRagCapability(query_engine)\n",
    "graph_rag_capability.add_to_agent(graph_rag_agent)\n",
    "\n",
    "# Create a user proxy agent to converse with our RAG agent\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy\",\n",
    "    human_input_mode=\"ALWAYS\",\n",
    ")\n",
    "\n",
    "user_proxy.initiate_chat(graph_rag_agent, message=\"Name a few actors who've played in 'The Matrix'\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Revisit the example with self defined graph ontology\n",
    "If you want to be more specific on the entities and attributes in the graph database you can create an `Ontology` before loading your data, in the case below we define these entities (`Actor` with `name` and `Movie` with `title`) and relationships (Actors `ACTED` in Movies). This allows the RAG agent to answer questions about actors in the movie."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from graphrag_sdk import Attribute, AttributeType, Entity, Ontology, Relation\n",
    "\n",
    "from autogen.agentchat.contrib.graph_rag.document import Document, DocumentType\n",
    "from autogen.agentchat.contrib.graph_rag.falkor_graph_query_engine import FalkorGraphQueryEngine\n",
    "\n",
    "input_path = \"../test/agentchat/contrib/graph_rag/the_matrix.txt\"\n",
    "\n",
    "movie_ontology = Ontology()\n",
    "movie_ontology.add_entity(\n",
    "    Entity(label=\"Actor\", attributes=[Attribute(name=\"name\", attr_type=AttributeType.STRING, unique=True)])\n",
    ")\n",
    "movie_ontology.add_entity(\n",
    "    Entity(label=\"Movie\", attributes=[Attribute(name=\"title\", attr_type=AttributeType.STRING, unique=True)])\n",
    ")\n",
    "movie_ontology.add_relation(Relation(label=\"ACTED\", source=\"Actor\", target=\"Movie\"))\n",
    "\n",
    "query_engine = FalkorGraphQueryEngine(\n",
    "    name=\"IMDB\",\n",
    "    host=\"172.17.0.4\",  # Change\n",
    "    port=6379,  # if needed\n",
    "    ontology=movie_ontology,\n",
    ")\n",
    "\n",
    "input_documents = [Document(doctype=DocumentType.TEXT, path_or_url=input_path)]\n",
    "\n",
    "query_engine.init_db(input_doc=input_documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Add capability to a ConversableAgent and query them\n",
    "\n",
    "With FalkorDB setup with an ontology and data, we can now associate the GraphRAG capability with a ConversableAgent and have a chat with it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen import ConversableAgent, UserProxyAgent\n",
    "from autogen.agentchat.contrib.graph_rag.falkor_graph_rag_capability import FalkorGraphRagCapability\n",
    "\n",
    "# Create a ConversableAgent (no LLM configuration)\n",
    "graph_rag_agent = ConversableAgent(\n",
    "    name=\"matrix_agent\",\n",
    "    human_input_mode=\"NEVER\",\n",
    ")\n",
    "\n",
    "# Associate the capability with the agent\n",
    "graph_rag_capability = FalkorGraphRagCapability(query_engine)\n",
    "graph_rag_capability.add_to_agent(graph_rag_agent)\n",
    "\n",
    "# Create a user proxy agent to converse with our RAG agent\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy\",\n",
    "    code_execution_config=False,\n",
    "    is_termination_msg=lambda msg: \"TERMINATE\" in msg[\"content\"],\n",
    "    human_input_mode=\"ALWAYS\",\n",
    ")\n",
    "\n",
    "user_proxy.initiate_chat(graph_rag_agent, message=\"Name a few actors who've played in 'The Matrix'\")\n",
    "\n",
    "# You will be prompted, as a human in the loop, after the response - feel free to ask more questions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the output we can see that:\n",
    "- We can ask about actors and movies because we defined the ontology (and they were also automatically created in the prior example)\n",
    "- There is inbuilt awareness of answers it has provided previously\n",
    "- It can't answer any questions outside of the actors and movies entities as we didn't include anything else in the ontology"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Remove graph from database:\n",
    "query_engine.delete()"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Using FalkorGraphRagCapability with agents for GraphRAG Question & Answering",
   "tags": [
    "RAG",
    "FalkorDB"
   ]
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
